import datetime

import arrow

from mongoengine import Q
from utilities.meta_view import Response, APIViewList, APIViewRetrieve
from models.vul import KbVuln
from apps.vul.serializer import VulListSchema, PieVulSchema, LineSchema
from django.core.paginator import EmptyPage
from collections import Counter

from apps.vul.serializer import PluginsListSerializer, plugin_filter


SEVERITY_LEVEL_MAP = {0: '未知', 1: '信息', 2: '低危', 3: '中危', 4: '高危', 5: '严重'}


def vul_unique():
    """漏洞中包含各种重复数据 需要做去重
    根据漏洞名称 去重 优先返回nscan数据
    """
    aggregate = [
        {'$sort': {'source': -1}},  # nscan > kscan > cve
        {'$group': {'_id': '$name.zh_cn', 'vid': {'$first': '$_id'}}}
    ]
    uniq_vids = [item['vid'] for item in KbVuln.objects.aggregate(aggregate)]
    return uniq_vids


class EnumType(APIViewList):
    def list(self):
        """枚举信息"""
        type_names = KbVuln.objects(**plugin_filter).distinct(
            'vuln_type'
        )  # 逻辑复用
        type_infos = [
            {'en_us': type_name.en_us, 'zh_cn': type_name.zh_cn}
            for type_name in type_names
        ]
        enum_info = {'severity': SEVERITY_LEVEL_MAP, 'vuln_type': type_infos}
        return Response(enum_info)


class ListVul(APIViewList):
    """漏洞接口 参照 scanv max 逻辑"""

    def list(self):
        q = VulListSchema(**self.request.GET.dict()).dict(
            skip_defaults=True, by_alias=True
        )
        # 需要特殊处理的查询字段
        start_time = q.pop('start_time', None)
        end_time = q.pop('end_time', None)
        keyword = q.pop('keyword', None)
        page = q.pop('page', 1)
        count = q.pop('count', 10)
        vuln_type = q.pop('vuln_type', None)

        if vuln_type:
            q['vuln_type.zh_cn'] = vuln_type

        updated_at = {}
        if start_time:
            updated_at['$gt'] = start_time
        if end_time:
            updated_at['$lt'] = end_time
        if updated_at:
            q['updated_at'] = updated_at

        if keyword:  # 该字段同时查名称/标签/编号  一个满足即可
            q1 = {'name.zh_cn': {'$regex': keyword, "$options": "i"}}
            q2 = {'tag': {'$regex': keyword, "$options": "i"}}
            q3 = {'cve': {'$regex': keyword, "$options": "i"}}  # todo 编号
            q['$or'] = [q1, q2, q3]

        q['_id'] = {'$in': vul_unique()}

        data = KbVuln.objects.find(q).order_by('-updated_at').paginate(count)
        total = data.count
        try:
            page_data = data.page(page)
        except EmptyPage:
            page_data = []

        new_data = []
        for item in page_data:
            new_data.append(
                {
                    'vuln_name': item.name.zh_cn,
                    'vuln_type': item.vuln_type[0].zh_cn if item.vuln_type else '',
                    'level': SEVERITY_LEVEL_MAP.get(item.severity, '未知'),
                    'published_at': arrow.get(item.published_at).isoformat()
                    if item.published_at
                    else '',
                    'updated_at': arrow.get(item.updated_at).isoformat()
                    if item.updated_at
                    else '',
                    'cnnvd': item.cnnvd,
                    'cve': item.cve,
                    'owasp': item.owasp,
                    'description': item.description.zh_cn,
                    'impact': item.impact.zh_cn if item.impact else '',
                    'recommendation': item.recommendation.zh_cn
                    if item.recommendation
                    else '',
                    'reference': item.reference,
                    'zoomeye_dork': item.zoomeye_dork,
                    'zoomeye_count': item.zoomeye_count,
                }
            )

        results = {'data': new_data, 'page': page, 'count': total}
        return Response({'code': 200, 'msg': 'ok', 'results': results})


class PieVul(APIViewList):
    """漏洞饼图"""

    def list(self):
        req_data = PieVulSchema(**self.request.GET.dict()).dict()

        pie_datas = {}
        vul_uids = vul_unique()
        for field in req_data['fields']:
            agg = [{'$match': {'_id': {'$in': vul_uids}}}, {'$group': {'_id': f'${field}', 'count': {'$sum': 1}}}]
            data = {
                item['_id']: item['count'] for item in KbVuln.objects.aggregate(agg)
            }
            pie_datas[field] = data
        return Response(pie_datas)


class LineVul(APIViewList):
    """线性图"""

    def list(self):
        default_data = (0, 0, 0, 0, 0)  # 总数 严重 高危 中 低
        now = datetime.datetime.now()
        req_data = LineSchema(**self.request.GET.dict()).dict()
        started = req_data['range']
        vul_uids = vul_unique()
        uniq_q = {'_id': {'$in': vul_uids}}

        q = {**uniq_q, 'created_at': {'$gt': started}}
        agg = [
            {'$match': q},
            {
                '$project': {
                    'created_at': {
                        '$dateToString': {
                            'format': "%Y-%m-%d",
                            'date': {"$add": ["$created_at", 8 * 3600000]},
                        }
                    },
                    'severity': 1,
                }
            },
            {'$group': {'_id': '$created_at', 'severity': {'$push': '$severity'}}},
            {'$sort': {'_id': 1}},
        ]
        date_map = {
            item['_id']: Counter(item['severity'])
            for item in KbVuln.objects.aggregate(agg)
        }
        data = []
        for day in range(0, (now - started).days + 2):
            date_str = (started + datetime.timedelta(days=day)).strftime('%Y-%m-%d')
            count = date_map.get(date_str)
            if count:
                data.append(
                    [
                        date_str,
                        *(
                            sum(count.values()),
                            count.get(5, 0),
                            count.get(4, 0),
                            count.get(3, 0),
                            count.get(2, 0),
                        ),
                    ]
                )
            else:
                data.append([date_str, *default_data])

        # 获取总数 近一周 近一月的增长数据量
        week_dt = now - datetime.timedelta(days=7)
        month_dt = now - datetime.timedelta(days=30)
        total = KbVuln.objects.find(uniq_q).count()
        weekly = KbVuln.objects.find({**uniq_q, 'created_at': {'$gte': week_dt}}).count()
        monthly = KbVuln.objects.find({**uniq_q, 'created_at': {'$gte': month_dt}}).count()

        datas = {
            'count': {'total': total, 'week': weekly, 'month': monthly},
            'lines': data,
        }
        return Response(datas)


class Vulst(APIViewRetrieve):
    """漏洞信息的各种极值
    最早 最晚 最大等信息
    """

    def get(self, *args, **kwargs):
        """当前取值 最新更新数据"""
        uniq_q = {'_id': {'$in': vul_unique()}}
        item = KbVuln.objects.find(uniq_q).order_by('-updated_at').first()
        if not item:
            data = {'updated_at': ''}
        else:
            data = {'updated_at': arrow.get(item.updated_at).isoformat()}
        return Response(data)


class SearchPluginsList(APIViewList):
    serializer_class = PluginsListSerializer
    queryset = PluginsListSerializer.Meta.model.objects(**plugin_filter)

    def list(self):
        search_key = self.request.GET.get('search_key')
        q = Q()
        if search_key:
            q = q & (Q(key__contains=search_key) | Q(name__zh_cn__contains=search_key))

        plugin_list = self.queryset.filter(q)
        serializer = self.get_serializer(plugin_list, many=True)
        return Response({'plugin_list': serializer.data})
